"""
@Time: 2020/12/28 下午 5:23
@Author: jinzhuan
@File: bert_ws.py
@Desc: 
"""
import sys
sys.path.append('/data/zhuoran/cognlp')
sys.path.append('/data/zhuoran/code/cognlp')

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import RandomSampler
from cognlp import *
from cognlp.io.loader.ws.msra import MSRALoader
from cognlp.io.processor.ws.msra import MsraWSProcessor


torch.cuda.set_device(2)
device = torch.device('cuda')
loader = MSRALoader()
train_data, dev_data, test_data = loader.load_all('../../../cognlp/data/ws/msra/data/')
processor = MsraWSProcessor(label_list=loader.get_labels(), path='../../../cognlp/data/ws/msra/data')
vocabulary = Vocabulary.load('../../../cognlp/data/ws/msra/data/vocabulary.txt')

# train_datable = processor.process(train_data)
# train_datable.save_table('../../../cognlp/data/ws/msra/data/train.json')
train_datable = DataTable.load_table('../../../cognlp/data/ws/msra/data/train.json')
train_dataset = DataTableSet(train_datable)
train_sampler = RandomSampler(train_dataset)

# dev_datable = processor.process(dev_data)
# dev_datable.save_table('../../../cognlp/data/ws/msra/data/dev.json')
dev_datable = DataTable.load_table('../../../cognlp/data/ws/msra/data/dev.json')
dev_dataset = DataTableSet(dev_datable)
dev_sampler = RandomSampler(dev_dataset)

# test_datable = processor.process(test_data)
# test_datable.save_table('../../../cognlp/data/ws/msra/data/test.json')
test_datable = DataTable.load_table('../../../cognlp/data/ws/msra/data/test.json')
test_dataset = DataTableSet(test_datable)
test_sampler = RandomSampler(test_dataset)

model = Bert4WS(vocabulary)
metric = SpanFPreRecMetric(vocabulary)
loss = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.00005)

trainer = Trainer(model,
                  train_dataset,
                  dev_data=test_dataset,
                  n_epochs=10,
                  batch_size=40,
                  loss=loss,
                  optimizer=optimizer,
                  scheduler=None,
                  metrics=metric,
                  train_sampler=train_sampler,
                  dev_sampler=test_sampler,
                  drop_last=False,
                  gradient_accumulation_steps=1,
                  num_workers=5,
                  save_path="../../../cognlp/data/ws/msra/model",
                  save_file=None,
                  print_every=None,
                  scheduler_steps=None,
                  validate_steps=None,
                  save_steps=None,
                  grad_norm=None,
                  use_tqdm=True,
                  device=device,
                  device_ids=[2, 3],
                  callbacks=None,
                  metric_key=None,
                  writer_path='../../../cognlp/data/ws/msra/tensorboard',
                  fp16=False,
                  fp16_opt_level='O1',
                  checkpoint_path=None,
                  task='msra-crf',
                  logger_path='../../../cognlp/data/ws/msra/logger')

trainer.train()
